import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import statsmodels.api as sm
from scipy import stats
from statsmodels.tsa.arima.model import ARIMA
from statsmodels.graphics.tsaplots import plot_predict
from sklearn.metrics import mean_squared_error
import matplotlib as mpl

mpl.use('TkAgg')

df = pd.read_csv("data/Q2.csv", index_col=0)
# 将df序号设置为date
df.index = df['date']
del df['date']

df = df.iloc[:, -1:]
print(df)

ARIMA201 = ARIMA(df["AQI"], order=(2, 0, 1) ).fit()  # 创建 ARIMA(2,0,1) 模型
ARIMA201.plot_diagnostics(figsize=(12, 8))
# plt.show()
plt.savefig("img/ARIMA201")



plot_predict(ARIMA201,start='2015',end='2023-4-29')
plt.savefig("img/Q3A_preAll.png")
# plt.show()

pred = ARIMA201.get_prediction(start='2015',end='2023-4-29')
pred_mean = pred.predicted_mean

plt.figure(figsize=(10, 6))
plt.plot(df.index, df["AQI"], label='observed')
plt.plot(df.index, pred_mean.values, label='forecast')
plt.gca().xaxis.set_visible(False)
# print(pred_mean.values)
# plt.fill_between(pred_ci.index, pred_ci.iloc[:, 0], pred_ci.iloc[:, 1], color='k', alpha=.2)
plt.legend()
plt.savefig("img/Q3A_bin.png")

pred_ = ARIMA201.forecast(steps=12)
print(pred_)

# pred1 = pd.DataFrame({"date":df.index,
#                       "pred":pred_mean.values})
# pred1.index = pred1["date"]
# del pred1["date"]
# print(pred1)
# pred1.to_csv("data/Q3A_pred1.csv")

